import argparse
import math
import numpy as np
import matplotlib.pyplot as plt
from numpy import genfromtxt

def box_size(ang, n, height, l):
    angle = math.radians(ang / n)  # 将夹角转换为弧度  
    radius = math.sqrt(height ** 2 + l ** 2)
    # 计算弧长  
    arc_length = angle * radius  
    return arc_length

def adjust_roll(roll_angle, x, y, z):
    pi = 3.1415926
    x = x * np.cos(roll_angle / 180 * pi) + z * np.sin(roll_angle / 180 * pi)
    y = y
    z = z * np.cos(roll_angle / 180 * pi) - x * np.sin(roll_angle / 180 * pi)
    return x, y, z

def adjust_pitch(pitch_angle, x, y, z):
    pi = 3.1415926
    x = x
    y = y * np.cos(pitch_angle / 180 * pi) + z * np.sin(pitch_angle / 180 * pi)
    z = z * np.cos(pitch_angle / 180 * pi) - y * np.sin(pitch_angle / 180 * pi)
    return x, y, z

def get_height(z):
    return z.mean()

def get_roll_angle(x_floor, y_floor, z_floor):
    maxangle = int(10)
    min_std = 999999
    roll_angle = 0

    for i in range(-maxangle, maxangle + 1, 1):
        theta1 = i
        pi = 3.1415926
        x1 = x_floor * np.cos(theta1 / 180 * pi) + z_floor * np.sin(theta1 / 180 * pi)
        y1 = y_floor
        z1 = z_floor * np.cos(theta1 / 180 * pi) - x_floor * np.sin(theta1 / 180 * pi)
        i_std = np.std(z1)
        if i_std < min_std:
            min_std = i_std
            roll_angle = i
    return roll_angle

def get_pitch_angle(x_floor, y_floor, z_floor):
    maxangle = int(10)
    min_std = 999999
    pitch_angle = 0

    for i in range(-maxangle, maxangle + 1, 1):
        theta2 = i
        pi = 3.1415926
        x2 = x_floor
        y2 = y_floor * np.cos(theta2 / 180 * pi) + z_floor * np.sin(theta2 / 180 * pi)
        z2 = z_floor * np.cos(theta2 / 180 * pi) - y_floor * np.sin(theta2 / 180 * pi)
        i_std = np.std(z2)
        if i_std < min_std:
            min_std = i_std
            pitch_angle = i
    return pitch_angle

def adjust_roll(roll_angle, x, y, z):
    pi = 3.1415926
    x = x * np.cos(roll_angle / 180 * pi) + z * np.sin(roll_angle / 180 * pi)
    y = y
    z = z * np.cos(roll_angle / 180 * pi) - x * np.sin(roll_angle / 180 * pi)
    return x, y, z

def adjust_pitch(pitch_angle, x, y, z):
    pi = 3.1415926
    x = x
    y = y * np.cos(pitch_angle / 180 * pi) + z * np.sin(pitch_angle / 180 * pi)
    z = z * np.cos(pitch_angle / 180 * pi) - y * np.sin(pitch_angle / 180 * pi)
    return x, y, z

def volume(grid_resolution, x_table, y_table, z_table, height, filename):
    x_range = np.arange(0, int(max(x_table)+1), grid_resolution)
    y_range = np.arange(0, int(max(y_table)+1), grid_resolution)
    x_range += int(0.5 * grid_resolution)
    y_range += int(0.5 * grid_resolution)   
    xx, yy = np.meshgrid(x_range, y_range)

    # 初始化颜色数组和计数器数组
    colors = np.zeros_like(xx)
    counts = np.zeros_like(xx)

    # 将点分配到网格并计算颜色
    for i in range(len(x_table)):
        x, y, z = x_table[i], y_table[i], z_table[i]
        x_index = int(x // grid_resolution)
        y_index = int(y // grid_resolution)
        colors[y_index, x_index] += z
        counts[y_index, x_index] += 1

    # 计算每个网格的平均z值
    colors = colors.astype(np.float64)
    colors /= np.where(counts == 0, 1, counts)
    # print(colors.sum())
    volume = 0
    area = grid_resolution * grid_resolution / 1e6
    for i in range(len(colors)):
        for j in range(len(colors[0])):
            if colors[i][j] != 0:
                area_height = (colors[i][j] - height) / 1000
                volume += area * area_height
    colors = [[x - height for x in sublist] for sublist in colors]
    # 创建网格图
    plt.figure(figsize=(30, 20))
    plt.pcolormesh(xx, yy, colors, shading='auto', cmap='viridis')
    plt.colorbar(label='Z')
    plt.scatter(x_table, y_table, s=20, c='red', marker='.')
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('The grid_resolution is %d. The volume of the table is %.5f' % (grid_resolution, volume))
    plt.gca().set_aspect('equal')  # 保持网格方格
    plt.savefig('./data/box/%s_%d.png' % (filename, grid_resolution), dpi=300, bbox_inches='tight')
    plt.show()
    plt.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser()  
    parser.add_argument("--n", default=1000, help="雷达采集一圈点的个数")  
    parser.add_argument("--angle", default=360)  
    parser.add_argument("--height", default=2700, help="雷达距离地面的高度")  
    parser.add_argument("--l", default=3000, help="雷达横向侧扫半径")  
    parser.add_argument("--floor_loc1", default=(-3000, 2500), help="用于校准数据的地面左上角坐标")  
    parser.add_argument("--floor_loc2", default=(1000, 2250), help="用于校准数据的地面右下角坐标")  
    parser.add_argument("--table_loc1", default=(-1000, 3500), help="测算物体所在区域的左上角坐标")  
    parser.add_argument("--table_loc2", default=(0, 2500), help="测算物体所在区域的右下角坐标")
    parser.add_argument("--file", default='C:/Users/赵浩中/PycharmProjects/pythonProject1/data/box/lidar-10231647.csv', help="读取文件路径") 
    args = parser.parse_args()

    d100 = genfromtxt(args.file, delimiter=',')
    x = -d100[:, 2]  # * np.cos(theta / 180 * pi) + d100[:, 2] * np.sin(theta / 180 * pi)
    z = -d100[:, 1]  # * np.cos(theta / 180 * pi) - d100[:, 1] * np.sin(theta / 180 * pi)
    y = d100[:, 3]
    # y = [i * 1.025 for i in d100[:, 3]]
    # y = np.array(y)
    condition_floor = (x > args.floor_loc1[0]) & (x < args.floor_loc2[0]) & (y > args.floor_loc2[1]) & (y < args.floor_loc1[1]) & (z < -2600)
    x_floor = x[condition_floor]        
    y_floor = y[condition_floor]
    z_floor = z[condition_floor]

    roll_angle = get_roll_angle(x_floor, y_floor, z_floor)
    x_floor, y_floor, z_floor = adjust_roll(roll_angle, x_floor, y_floor, z_floor)
    pitch_angle = get_pitch_angle(x_floor, y_floor, z_floor)
    x_floor, y_floor, z_floor = adjust_pitch(pitch_angle, x_floor, y_floor, z_floor)
    height = get_height(z_floor)

    condition_table = (x > args.table_loc1[0]) & (x < args.table_loc2[0]) & (y > args.table_loc2[1]) & (y < args.table_loc1[1]) & (z < -1500)
    x_table = x[condition_table]
    y_table = y[condition_table]
    z_table = z[condition_table]
    x_table, y_table, z_table = adjust_roll(roll_angle, x_table, y_table, z_table)
    x_table, y_table, z_table = adjust_pitch(pitch_angle, x_table, y_table, z_table)
    x_table_min, x_table_max = min(x_table), max(x_table)
    y_table_min, y_table_max = min(y_table), max(y_table)
    for i in range(len(x_table)):
        if x_table_min < 0:
            x_table[i] += abs(x_table_min)
        else:
            x_table[i] -= abs(x_table_min)
        y_table[i] -= abs(y_table_min)

    grid_resolution = int(box_size(args.angle, args.n, height, args.l))
    # print(height, grid_resolution)
    # exit()
    filename = args.file.split('/')[-1].split('.')[0]
    volume(grid_resolution, x_table, y_table, z_table, height, filename)